% main_twostates
% Jones, Philippon, Venkateswaran (2021)
%

clear 
clc

%%

warning('Off','all') ;

path(pathdef) ;

%% Deterministic, No Fatigue

planner = 1 ;

e0exog   = 0 ;
Imax     = 0.2;              %% targets for mortality
BaseMort = 0.0025;           
PeakMort = 3*BaseMort;
deterministic_vac = 1 ;
stoch_prob = 77/78 ;

beta       = (0.97)^(1/52)*stoch_prob ; 
ee_T_vac   = zeros(1,1000); 
ee_T_e0    = zeros(1,1000); 
ee_T_d     = zeros(1,1000); 
if deterministic_vac; beta = (0.97)^(1/52); T_vac = 1*52 ; ee_T_vac(T_vac:end)=0.01; for ii=(T_vac+1):length(ee_T_vac); ee_T_vac(ii)=ee_T_vac(ii-1)*1.2; end ; ee_T_vac(ee_T_vac>1) = 1; end
e0         = e0exog        + (1-e0exog)*(1/2) ;
ec1        = e0exog*0.0001 + (1-e0exog)*((1-e0)/2) ;
el1        = e0exog*0.0001 + (1-e0exog)*((1-e0)/2) ;
R0         = 2.0 ;
chibar1    = 0.5 ;
kappa      = 0.15;                                    % Fraction severely sick
rho        = 0.35;                                    % Recovery rate of infected 
deltabar   = BaseMort/(1-BaseMort)*rho/kappa;         % baseline death rate 
phi        = log(1-deltabar+PeakMort/(1-PeakMort)*rho/kappa)/(Imax*kappa);
N          = 1 ;
eta        = 2 ;
rhom       = 1  ;
Deltachi   = 0.34 ;
us         = 0.5 ;
ud         = 2.5 ;
ee_i0_init = (1/100)*N ; ee_T_i = zeros(1,1000) ; ee_T_i(1) = ee_i0_init ;
psim       = 0.99 ;
deltascale = 1 ;

rhof        = 0.0 ;
omega1      = 0.0 ;
omega2      = 2 ;

rhoa       = 0.0 ;
varphi1    = 0.5 ;
varphi2    = 0.2 ;
amean      = 1 ;

if planner
    model_name = 'infections_planner' ;
    set_params_planner
else
    model_name = 'infections' ;
    set_params
end

run_dynare

figure(1)

set(gcf, 'DefaultAxesLineWidth', 1.5);
set(gcf, 'DefaultLineLineWidth', 2);
set(gcf, 'DefaultAxesTickLabelInterpreter','latex'); 
set(gcf, 'DefaultLegendInterpreter','latex');
set(gcf, 'DefaultAxesFontSize',10);

subplot(2,2,2); hold on;
h1=plot(2:(length(s)-1),e(3:end)) ;
line([0 300],[1 1],'LineWidth',1,'Color','k') ;
title('Exposure, $e$','Interpreter','latex')
xlim([0 100]);

subplot(2,2,1); 
yyaxis left; hold on ;
h1 = plot(0:100,C1(1:101),'-','Color',h1.Color) ;
title('Consumption (LHS), Labor (RHS)','Interpreter','latex')
xlim([0 100]);
set(gca,'YColor','k') ;
ylim([0.2 1.2]) ;

yyaxis right; hold on ;
h1 = plot(0:100,L1(1:101),'--','Color',h1.Color) ;
title('Consumption (LHS), Labor (RHS)','Interpreter','latex')
xlim([0 100]);
set(gca,'YColor','k') ;
ylim([0.2 1.2]) ;

subplot(2,2,3); hold on;
h1 = plot(0:(length(s)-1),100*in) ;
title('Infected, \%','Interpreter','latex')
xlim([0 100]);

subplot(2,2,4); hold on;
yyaxis left
hold on ;
plot(0:(length(s)-1),100*s,'-','Color',h1.Color) ;
xlim([0 100]);
set(gca,'YColor','k') ;

yyaxis right
hold on ;
plot(0:(length(d)-1),100*d,'--','Color',h1.Color) ;
xlim([0 100]);
set(gca,'YColor','k') ;

title('Susceptible (LHS), Dead (RHS), \%','Interpreter','latex')

set(findall(gcf, 'Type', 'Text'),'FontWeight', 'Normal')

disp('Mean Decline in Consumption at 26 Weeks');
mean(C1(2:27))
disp('Mean Decline in Labor at 26 Weeks');
mean(L1(2:27))

disp('Deaths at 40 Weeks');
mean(d(41)*100)

disp('Deaths at end');
mean(d(200)*100)

%% 

planner = 1 ;

e0exog   = 0 ;
Imax     = 0.2;              %% targets for mortality
BaseMort = 0.0025;           
PeakMort = 3*BaseMort;
deterministic_vac = 1 ;
stoch_prob = 77/78 ;

beta       = (0.97)^(1/52)*stoch_prob ; 
ee_T_vac   = zeros(1,1000); 
ee_T_e0    = zeros(1,1000); 
ee_T_d     = zeros(1,1000); 
if deterministic_vac; beta = (0.97)^(1/52); T_vac = 1*52 ; ee_T_vac(T_vac:end)=0.01; for ii=(T_vac+1):length(ee_T_vac); ee_T_vac(ii)=ee_T_vac(ii-1)*1.2; end ; ee_T_vac(ee_T_vac>1) = 1; end
e0         = e0exog        + (1-e0exog)*(1/2) ;
ec1        = e0exog*0.0001 + (1-e0exog)*((1-e0)/2) ;
el1        = e0exog*0.0001 + (1-e0exog)*((1-e0)/2) ;
R0         = 2.0 ;
chibar1    = 0.5 ;
kappa      = 0.15;                                    % Fraction severely sick
rho        = 0.35;                                    % Recovery rate of infected 
deltabar   = BaseMort/(1-BaseMort)*rho/kappa;         % baseline death rate 
phi        = log(1-deltabar+PeakMort/(1-PeakMort)*rho/kappa)/(Imax*kappa);
N          = 1 ;
eta        = 2 ;
rhom       = 1  ;
Deltachi   = 0.34 ;
us         = 0.5 ;
ud         = 2.5 ;
ee_i0_init = (1/10000)*N ; ee_T_i = zeros(1,1000) ; ee_T_i(1) = ee_i0_init ;
psim       = 0.99 ;
deltascale = 1 ;

rhof        = 0.0 ;
omega1      = 0.0 ;
omega2      = 2 ;

rhoa       = 0.0 ;
varphi1    = 0.5 ;
varphi2    = 0.2 ;
amean      = 1 ;

if planner
    model_name = 'infections_planner' ;
    set_params_planner
else
    model_name = 'infections' ;
    set_params
end

run_dynare


figure(1)

set(gcf, 'DefaultAxesLineWidth', 1.5);
set(gcf, 'DefaultLineLineWidth', 2);
set(gcf, 'DefaultAxesTickLabelInterpreter','latex'); 
set(gcf, 'DefaultLegendInterpreter','latex');
set(gcf, 'DefaultAxesFontSize',10);

subplot(2,2,2); hold on;
h1=plot(2:(length(s)-1),e(3:end)) ;
line([0 300],[1 1],'LineWidth',1,'Color','k') ;
title('Exposure, $e$','Interpreter','latex')
xlim([0 75]);

subplot(2,2,1); 
yyaxis left; hold on ;
h1 = plot(0:100,C1(1:101),'-','Color',h1.Color) ;
title('Consumption (LHS), Labor (RHS)','Interpreter','latex')
xlim([0 75]);
set(gca,'YColor','k') ;
ylim([0.75 1.05]) ;

yyaxis right; hold on ;
h1 = plot(0:100,L1(1:101),'--','Color',h1.Color) ;
title('Cons. (Solid), Labor (Dashed)','Interpreter','latex')
xlim([0 75]);
set(gca,'YColor','k') ;
ylim([0.75 1.05]) ;

subplot(2,2,3); hold on;
h1 = plot(0:(length(s)-1),100*in) ;
title('Infected, \%','Interpreter','latex')
xlim([0 75]);

subplot(2,2,4); hold on;
yyaxis left
hold on ;
plot(0:(length(s)-1),100*s,'-','Color',h1.Color) ;
xlim([0 75]);
set(gca,'YColor','k') ;

yyaxis right
hold on ;
plot(0:(length(d)-1),100*d,'--','Color',h1.Color) ;
xlim([0 75]);
set(gca,'YColor','k') ;

title('Susceptible (LHS), Dead (RHS), \%','Interpreter','latex')

set(findall(gcf, 'Type', 'Text'),'FontWeight', 'Normal')

print -depsc ./output/twostates.eps
print -dpng ./output/twostates.png
close

disp('Mean Decline in Consumption at 26 Weeks');
mean(C1(2:27))
disp('Mean Decline in Labor at 26 Weeks');
mean(L1(2:27))

disp('Deaths at 40 Weeks');
mean(d(41)*100)

disp('Deaths at end');
mean(d(200)*100)
